#include "md5_4.h"
#include <iomanip>
#include <assert.h>
#include <chrono>
#include <omp.h>

using namespace std;
using namespace chrono;

/**
 * StringProcess_SSE: 将单个输入字符串转换成MD5计算所需的消息数组
 * @param input 输入
 * @param[out] n_byte 用于给调用者传递额外的返回值，即最终Byte数组的长度
 * @return Byte消息数组
 */
Byte *StringProcess_SSE(string input, int *n_byte)
{
    // 将输入的字符串转换为Byte为单位的数组
    Byte *blocks = (Byte *)input.c_str();
    int length = input.length();

    // 计算原始消息长度（以比特为单位）
    int bitLength = length * 8;

    // paddingBits: 原始消息需要的padding长度（以bit为单位）
    // 对于给定的消息，将其补齐至length%512==448为止
    // 需要注意的是，即便给定的消息满足length%512==448，也需要再pad 512bits
    int paddingBits = bitLength % 512;
    if (paddingBits > 448)
    {
        paddingBits += 512 - (paddingBits - 448);
    }
    else if (paddingBits < 448)
    {
        paddingBits = 448 - paddingBits;
    }
    else if (paddingBits == 448)
    {
        paddingBits = 512;
    }

    // 原始消息需要的padding长度（以Byte为单位）
    int paddingBytes = paddingBits / 8;
    // 创建最终的字节数组
    // length + paddingBytes + 8:
    // 1. length为原始消息的长度（bits）
    // 2. paddingBytes为原始消息需要的padding长度（Bytes）
    // 3. 在pad到length%512==448之后，需要额外附加64bits的原始消息长度，即8个bytes
    int paddedLength = length + paddingBytes + 8;
    Byte *paddedMessage = new Byte[paddedLength];

    // 复制原始消息
    memcpy(paddedMessage, blocks, length);

    // 添加填充字节。填充时，第一位为1，后面的所有位均为0。
    // 所以第一个byte是0x80
    paddedMessage[length] = 0x80;                             // 添加一个0x80字节
    memset(paddedMessage + length + 1, 0, paddingBytes - 1); // 填充0字节

    // 添加消息长度（64比特，小端格式）
    for (int i = 0; i < 8; ++i)
    {
        // 特别注意此处应当将bitLength转换为uint64_t
        // 这里的length是原始消息的长度
        paddedMessage[length + paddingBytes + i] = ((uint64_t)length * 8 >> (i * 8)) & 0xFF;
    }

    // 验证长度是否满足要求。此时长度应当是512bit的倍数
    int residual = 8 * paddedLength % 512;
    // assert(residual == 0);

    // 在填充+添加长度之后，消息被分为n_blocks个512bit的部分
    *n_byte = paddedLength;
    return paddedMessage;
}

/**
 * SSE_MD5Hash_4_SSE: SSE指令集实现的四路并行MD5哈希计算
 * @param input 输入的4个字符串
 * @param[out] state 用于存储4个MD5哈希结果
 */
void SSE_MD5Hash_4_SSE(string input[4], bit32 state[4][4])
{
    Byte *paddedMessage[4];
    int *messageLength = new int[4];
    
    // 处理4个输入字符串
    for (int i = 0; i < 4; i++)
    {
        paddedMessage[i] = StringProcess_SSE(input[i], &messageLength[i]);
    }
    
    // 计算每个输入字符串的块数
    int n_blocks[4] = {
        messageLength[0] / 64, 
        messageLength[1] / 64, 
        messageLength[2] / 64, 
        messageLength[3] / 64
    };
    
    // 找出最大和最小块数
    int max_n_blocks = n_blocks[0];
    int min_n_blocks = n_blocks[0];
    for (int i = 1; i < 4; i++)
    {
        if (n_blocks[i] > max_n_blocks)
        {
            max_n_blocks = n_blocks[i];
        }
        if (n_blocks[i] < min_n_blocks)
        {
            min_n_blocks = n_blocks[i];
        }
    }

    // 初始化MD5状态
    for (int i = 0; i < 4; i++)
    {
        state[i][0] = 0x67452301;
        state[i][1] = 0xefcdab89;
        state[i][2] = 0x98badcfe;
        state[i][3] = 0x10325476;
    }
    
    // 创建SSE向量类型的变量
    __m128i a_end = _mm_set_epi32(state[3][0], state[2][0], state[1][0], state[0][0]);
    __m128i b_end = _mm_set_epi32(state[3][1], state[2][1], state[1][1], state[0][1]);
    __m128i c_end = _mm_set_epi32(state[3][2], state[2][2], state[1][2], state[0][2]);
    __m128i d_end = _mm_set_epi32(state[3][3], state[2][3], state[1][3], state[0][3]);
    
    __m128i a = a_end;
    __m128i b = b_end;
    __m128i c = c_end;
    __m128i d = d_end;
    __m128i x[16];
    
    int over_num = 0;
    
    // 逐块处理
    for (int i = 0; i <= max_n_blocks; i++)
    {
        // 检查是否完成处理每个输入
        if (i >= min_n_blocks)
        {
            // 检查每个输入是否处理完毕
            if (i == n_blocks[0])
            {
                state[0][0] = _mm_extract_epi32(a_end, 0);
                state[0][1] = _mm_extract_epi32(b_end, 0);
                state[0][2] = _mm_extract_epi32(c_end, 0);
                state[0][3] = _mm_extract_epi32(d_end, 0);
                over_num++;
            }
            if (i == n_blocks[1])
            {
                state[1][0] = _mm_extract_epi32(a_end, 1);
                state[1][1] = _mm_extract_epi32(b_end, 1);
                state[1][2] = _mm_extract_epi32(c_end, 1);
                state[1][3] = _mm_extract_epi32(d_end, 1);
                over_num++;
            }
            if (i == n_blocks[2])
            {
                state[2][0] = _mm_extract_epi32(a_end, 2);
                state[2][1] = _mm_extract_epi32(b_end, 2);
                state[2][2] = _mm_extract_epi32(c_end, 2);
                state[2][3] = _mm_extract_epi32(d_end, 2);
                over_num++;
            }
            if (i == n_blocks[3])
            {
                state[3][0] = _mm_extract_epi32(a_end, 3);
                state[3][1] = _mm_extract_epi32(b_end, 3);
                state[3][2] = _mm_extract_epi32(c_end, 3);
                state[3][3] = _mm_extract_epi32(d_end, 3);
                over_num++;
            }
            if (over_num == 4)
            {
                break;
            }
        }
        
        // 加载数据到SIMD寄存器
        for (int j = 0; j < 16; j++)
        {
            int offset = 4 * j + i * 64;
            uint32_t val0 = *reinterpret_cast<uint32_t*>(paddedMessage[0] + offset);
            uint32_t val1 = *reinterpret_cast<uint32_t*>(paddedMessage[1] + offset);
            uint32_t val2 = *reinterpret_cast<uint32_t*>(paddedMessage[2] + offset);
            uint32_t val3 = *reinterpret_cast<uint32_t*>(paddedMessage[3] + offset);
            
            x[j] = _mm_set_epi32(val3, val2, val1, val0);
        }
        
        // 保存当前状态
        a = a_end;
        b = b_end;
        c = c_end;
        d = d_end;
        
        /* Round 1 */
        SSE_FF(a, b, c, d, x[0], s11, 0xd76aa478);
        SSE_FF(d, a, b, c, x[1], s12, 0xe8c7b756);
        SSE_FF(c, d, a, b, x[2], s13, 0x242070db);
        SSE_FF(b, c, d, a, x[3], s14, 0xc1bdceee);
        SSE_FF(a, b, c, d, x[4], s11, 0xf57c0faf);
        SSE_FF(d, a, b, c, x[5], s12, 0x4787c62a);
        SSE_FF(c, d, a, b, x[6], s13, 0xa8304613);
        SSE_FF(b, c, d, a, x[7], s14, 0xfd469501);
        SSE_FF(a, b, c, d, x[8], s11, 0x698098d8);
        SSE_FF(d, a, b, c, x[9], s12, 0x8b44f7af);
        SSE_FF(c, d, a, b, x[10], s13, 0xffff5bb1);
        SSE_FF(b, c, d, a, x[11], s14, 0x895cd7be);
        SSE_FF(a, b, c, d, x[12], s11, 0x6b901122);
        SSE_FF(d, a, b, c, x[13], s12, 0xfd987193);
        SSE_FF(c, d, a, b, x[14], s13, 0xa679438e);
        SSE_FF(b, c, d, a, x[15], s14, 0x49b40821);

        /* Round 2 */
        SSE_GG(a, b, c, d, x[1], s21, 0xf61e2562);
        SSE_GG(d, a, b, c, x[6], s22, 0xc040b340);
        SSE_GG(c, d, a, b, x[11], s23, 0x265e5a51);
        SSE_GG(b, c, d, a, x[0], s24, 0xe9b6c7aa);
        SSE_GG(a, b, c, d, x[5], s21, 0xd62f105d);
        SSE_GG(d, a, b, c, x[10], s22, 0x2441453);
        SSE_GG(c, d, a, b, x[15], s23, 0xd8a1e681);
        SSE_GG(b, c, d, a, x[4], s24, 0xe7d3fbc8);
        SSE_GG(a, b, c, d, x[9], s21, 0x21e1cde6);
        SSE_GG(d, a, b, c, x[14], s22, 0xc33707d6);
        SSE_GG(c, d, a, b, x[3], s23, 0xf4d50d87);
        SSE_GG(b, c, d, a, x[8], s24, 0x455a14ed);
        SSE_GG(a, b, c, d, x[13], s21, 0xa9e3e905);
        SSE_GG(d, a, b, c, x[2], s22, 0xfcefa3f8);
        SSE_GG(c, d, a, b, x[7], s23, 0x676f02d9);
        SSE_GG(b, c, d, a, x[12], s24, 0x8d2a4c8a);

        /* Round 3 */
        SSE_HH(a, b, c, d, x[5], s31, 0xfffa3942);
        SSE_HH(d, a, b, c, x[8], s32, 0x8771f681);
        SSE_HH(c, d, a, b, x[11], s33, 0x6d9d6122);
        SSE_HH(b, c, d, a, x[14], s34, 0xfde5380c);
        SSE_HH(a, b, c, d, x[1], s31, 0xa4beea44);
        SSE_HH(d, a, b, c, x[4], s32, 0x4bdecfa9);
        SSE_HH(c, d, a, b, x[7], s33, 0xf6bb4b60);
        SSE_HH(b, c, d, a, x[10], s34, 0xbebfbc70);
        SSE_HH(a, b, c, d, x[13], s31, 0x289b7ec6);
        SSE_HH(d, a, b, c, x[0], s32, 0xeaa127fa);
        SSE_HH(c, d, a, b, x[3], s33, 0xd4ef3085);
        SSE_HH(b, c, d, a, x[6], s34, 0x4881d05);
        SSE_HH(a, b, c, d, x[9], s31, 0xd9d4d039);
        SSE_HH(d, a, b, c, x[12], s32, 0xe6db99e5);
        SSE_HH(c, d, a, b, x[15], s33, 0x1fa27cf8);
        SSE_HH(b, c, d, a, x[2], s34, 0xc4ac5665);

        /* Round 4 */
        SSE_II(a, b, c, d, x[0], s41, 0xf4292244);
        SSE_II(d, a, b, c, x[7], s42, 0x432aff97);
        SSE_II(c, d, a, b, x[14], s43, 0xab9423a7);
        SSE_II(b, c, d, a, x[5], s44, 0xfc93a039);
        SSE_II(a, b, c, d, x[12], s41, 0x655b59c3);
        SSE_II(d, a, b, c, x[3], s42, 0x8f0ccc92);
        SSE_II(c, d, a, b, x[10], s43, 0xffeff47d);
        SSE_II(b, c, d, a, x[1], s44, 0x85845dd1);
        SSE_II(a, b, c, d, x[8], s41, 0x6fa87e4f);
        SSE_II(d, a, b, c, x[15], s42, 0xfe2ce6e0);
        SSE_II(c, d, a, b, x[6], s43, 0xa3014314);
        SSE_II(b, c, d, a, x[13], s44, 0x4e0811a1);
        SSE_II(a, b, c, d, x[4], s41, 0xf7537e82);
        SSE_II(d, a, b, c, x[11], s42, 0xbd3af235);
        SSE_II(c, d, a, b, x[2], s43, 0x2ad7d2bb);
        SSE_II(b, c, d, a, x[9], s44, 0xeb86d391);

        // 累加结果
        a_end = _mm_add_epi32(a_end, a);
        b_end = _mm_add_epi32(b_end, b);
        c_end = _mm_add_epi32(c_end, c);
        d_end = _mm_add_epi32(d_end, d);
    }
    
    // 处理最终的状态值，将字节序从小端转换为正确的格式
    // 修复: 展开循环，使用常量索引代替变量索引
    // j=0 的情况
    {
        uint32_t a_val = _mm_extract_epi32(a_end, 0);
        uint32_t b_val = _mm_extract_epi32(b_end, 0);
        uint32_t c_val = _mm_extract_epi32(c_end, 0);
        uint32_t d_val = _mm_extract_epi32(d_end, 0);
        
        // 字节序转换
        state[0][0] = ((a_val & 0xff) << 24) | ((a_val & 0xff00) << 8) | 
                      ((a_val & 0xff0000) >> 8) | ((a_val & 0xff000000) >> 24);
        state[0][1] = ((b_val & 0xff) << 24) | ((b_val & 0xff00) << 8) | 
                      ((b_val & 0xff0000) >> 8) | ((b_val & 0xff000000) >> 24);
        state[0][2] = ((c_val & 0xff) << 24) | ((c_val & 0xff00) << 8) | 
                      ((c_val & 0xff0000) >> 8) | ((c_val & 0xff000000) >> 24);
        state[0][3] = ((d_val & 0xff) << 24) | ((d_val & 0xff00) << 8) | 
                      ((d_val & 0xff0000) >> 8) | ((d_val & 0xff000000) >> 24);
    }
    
    // j=1 的情况
    {
        uint32_t a_val = _mm_extract_epi32(a_end, 1);
        uint32_t b_val = _mm_extract_epi32(b_end, 1);
        uint32_t c_val = _mm_extract_epi32(c_end, 1);
        uint32_t d_val = _mm_extract_epi32(d_end, 1);
        
        // 字节序转换
        state[1][0] = ((a_val & 0xff) << 24) | ((a_val & 0xff00) << 8) | 
                      ((a_val & 0xff0000) >> 8) | ((a_val & 0xff000000) >> 24);
        state[1][1] = ((b_val & 0xff) << 24) | ((b_val & 0xff00) << 8) | 
                      ((b_val & 0xff0000) >> 8) | ((b_val & 0xff000000) >> 24);
        state[1][2] = ((c_val & 0xff) << 24) | ((c_val & 0xff00) << 8) | 
                      ((c_val & 0xff0000) >> 8) | ((c_val & 0xff000000) >> 24);
        state[1][3] = ((d_val & 0xff) << 24) | ((d_val & 0xff00) << 8) | 
                      ((d_val & 0xff0000) >> 8) | ((d_val & 0xff000000) >> 24);
    }
    
    // j=2 的情况
    {
        uint32_t a_val = _mm_extract_epi32(a_end, 2);
        uint32_t b_val = _mm_extract_epi32(b_end, 2);
        uint32_t c_val = _mm_extract_epi32(c_end, 2);
        uint32_t d_val = _mm_extract_epi32(d_end, 2);
        
        // 字节序转换
        state[2][0] = ((a_val & 0xff) << 24) | ((a_val & 0xff00) << 8) | 
                      ((a_val & 0xff0000) >> 8) | ((a_val & 0xff000000) >> 24);
        state[2][1] = ((b_val & 0xff) << 24) | ((b_val & 0xff00) << 8) | 
                      ((b_val & 0xff0000) >> 8) | ((b_val & 0xff000000) >> 24);
        state[2][2] = ((c_val & 0xff) << 24) | ((c_val & 0xff00) << 8) | 
                      ((c_val & 0xff0000) >> 8) | ((c_val & 0xff000000) >> 24);
        state[2][3] = ((d_val & 0xff) << 24) | ((d_val & 0xff00) << 8) | 
                      ((d_val & 0xff0000) >> 8) | ((d_val & 0xff000000) >> 24);
    }
    
    // j=3 的情况
    {
        uint32_t a_val = _mm_extract_epi32(a_end, 3);
        uint32_t b_val = _mm_extract_epi32(b_end, 3);
        uint32_t c_val = _mm_extract_epi32(c_end, 3);
        uint32_t d_val = _mm_extract_epi32(d_end, 3);
        
        // 字节序转换
        state[3][0] = ((a_val & 0xff) << 24) | ((a_val & 0xff00) << 8) | 
                      ((a_val & 0xff0000) >> 8) | ((a_val & 0xff000000) >> 24);
        state[3][1] = ((b_val & 0xff) << 24) | ((b_val & 0xff00) << 8) | 
                      ((b_val & 0xff0000) >> 8) | ((b_val & 0xff000000) >> 24);
        state[3][2] = ((c_val & 0xff) << 24) | ((c_val & 0xff00) << 8) | 
                      ((c_val & 0xff0000) >> 8) | ((c_val & 0xff000000) >> 24);
        state[3][3] = ((d_val & 0xff) << 24) | ((d_val & 0xff00) << 8) | 
                      ((d_val & 0xff0000) >> 8) | ((d_val & 0xff000000) >> 24);
    }
    
    // 释放内存
    for (int i = 0; i < 4; i++)
    {
        delete[] paddedMessage[i];
    }
    delete[] messageLength;
}

/**
 * SSE_MD5Hash_8_SSE: 使用OpenMP实现八路并行，本质是调用两次四路并行
 * @param input 输入的8个字符串
 * @param[out] state 用于存储8个MD5哈希结果
 */
void SSE_MD5Hash_8_SSE(string input[8], bit32 state[8][4])
{
#pragma omp parallel sections
    {
#pragma omp section
        {
            SSE_MD5Hash_4_SSE(input, state); // 处理前4个口令
        }
#pragma omp section
        {
            SSE_MD5Hash_4_SSE(input + 4, state + 4); // 处理后4个口令
        }
    }
}
